{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "636497e9-825c-43b1-b2ee-dc483c110a70",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.append(\"../..\")  # make it possible to import from the experimental dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dce46c9-2a8e-43f8-ad3f-d410ba1445f7",
   "metadata": {},
   "source": [
    "ColBERT is currently an experimental feature. ColBERT is a neural retrieval method that tends to work well in a zero-shot setting on out of domain datasets, due to it's use of token-level encodings (rather than at the sentence or chunk level)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4a52cb17-6cac-41e0-8d74-1d4ef03ee6e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import SimpleDirectoryReader, ServiceContext\n",
    "from experimental.colbert_index import ColbertIndex\n",
    "from llama_index.llms import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "10cb46a9-1cd6-4eff-9aee-34861aa43bef",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n",
    "service_context = ServiceContext.from_defaults(llm=llm, embed_model=\"local\")\n",
    "\n",
    "documents = SimpleDirectoryReader(\"data\").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ffaa8b5-2d82-4544-a54d-efe27fcce538",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = ColbertIndex.from_documents(\n",
    "    documents=documents, service_context=service_context\n",
    ")\n",
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\"What did the author do after his time at Y Combinator?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5f483a87-e5a0-456f-9b71-1b747a256627",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "After his time at Y Combinator, the author moved back to Providence to continue his studies at RISD. He then moved to New York City and took up residence in Yorkville, where he began to make paintings and pursue his career as an artist.\n"
     ]
    }
   ],
   "source": [
    "print(response.response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c02ec04-1f7d-4d60-a9f9-742a32c9ec97",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
